AbstractField

@beartype
class AbstractField(ABC):
    """Abstract base class for fields on a mesh.

    Defines the interface that ``ScalarField`` and ``VectorField`` implement.
    Holds a reference to the underlying mesh, the finite element type, and the
    degree-of-freedom values.

    Attributes:
        mesh: The mesh the field is defined on.
        element_type: The finite element type used for the field.
        value: The degree-of-freedom values of the field.

    """

    mesh: Mesh

    def __init__(
        self: Self,
        mesh: Mesh,
        element_type: skfem.element.Element,
        value: Number | list | tuple | NDArray | Callable,
        match_quadrature: Self | None = None,
    ):
        """Initialize the field.

        Args:
            mesh: The mesh the field is defined on.
            element_type: The scikit-fem element used for the field.
            value: Initial field values at the degrees of freedom. Can be a
                scalar (constant field), an array, or a callable that maps DOF
                coordinates to values.
            match_quadrature: Another field whose quadrature scheme should be
                reused for this field. Defaults to None.

        """
        self._match_quadrature = match_quadrature
        self.mesh = mesh
        self._element_type = element_type
        self.value = value

    @property
    @abstractmethod
    def dof_locations(self: Self) -> NDArray:
        """Return the coordinates of all degrees of freedom.

        Returns:
            Array of DOF coordinates with shape ``(n_dofs, 3)``.

        """
        pass

    @abstractmethod
    def get_subregion_dof_index(self: Self, name: str) -> NDArray:
        """Return DOF indices belonging to a named subregion or field boundary.

        Args:
            name: Name of the subregion.

        Returns:
            Array of DOF indices.

        Raises:
            ValueError: If the subregion is not found in either subdomains or
                boundaries.

        """
        pass

    @property
    @abstractmethod
    def element_type(self: Self) -> skfem.Element:
        """Return the finite element type.

        Returns:
            A scikit-fem element instance.

        """
        pass

    @property
    @abstractmethod
    def value(self: Self) -> NDArray:
        """Return the field values at the degrees of freedom.

        Returns:
            Array of field values.

        """
        pass

    @value.setter
    @abstractmethod
    def value(self: Self, value: Number | list | tuple | NDArray | Callable):
        pass

    @abstractmethod
    def set_subregion_value(
        self: Self, name: str, value: Number | list | tuple | NDArray | Callable
    ):
        """Set field values on a named subregion.

        Args:
            name: Name of the subregion to set values on.
            value: Value(s) to assign. Accepts the same types as the
                ``value`` setter.

        """
        pass

    @property
    def skfem_basis(self: Self) -> skfem.Basis:
        """Return the scikit-fem basis for this field.

        Returns:
            A scikit-fem ``Basis`` object constructed from the mesh and element type.

        Notes:
            If a ``match_quadrature`` field was provided, the basis reuses that
            field's quadrature scheme instead of creating a new one.

        """
        if self._match_quadrature is None:
            return skfem.Basis(self.mesh.to_skfem(), self.element_type)
        else:
            if self.mesh == self._match_quadrature.mesh:
                return self._match_quadrature.skfem_basis.with_element(
                    self.element_type
                )
            else:
                raise RuntimeError(
                    f"Cannot set quadrature from {self._match_quadrature} "
                    f"since the mesh {self._match_quadrature.mesh} doesnot match with "
                    f"{self.mesh}."
                )

    @abstractmethod
    def to_skfem(self: Self) -> skfem.DiscreteField:
        """Convert the field to a scikit-fem discrete field.

        Returns:
            A scikit-fem ``DiscreteField`` instance.

        """
        pass

    @abstractmethod
    def to_pyvista(self: Self) -> pv.UnstructuredGrid:
        """Convert the field to a PyVista UnstructuredGrid object.

        Returns:
            PyVista ``UnstructuredGrid`` with point data representing the field.

        """
        pass

    def to_file(self: Self, file_path: pl.Path | str):
        """Write the field to disk.

        Args:
            file_path: Output path. Supported extensions are ``.h5`` (HDF5)
                and ``.vtu`` (VTK UnstructuredGrid).

        Raises:
            NotImplementedError: If the file extension is not supported.

        """
        file_path = file_path if isinstance(file_path, pl.Path) else pl.Path(file_path)
        if file_path.suffix == ".h5":
            self.mesh.to_file(file_path)
            with h5py.File(name=file_path, mode="r+", track_order=True) as handle:
                field = handle.create_group("field")
                field_type = type(self).__name__
                field.attrs["type"] = field_type
                field.attrs["element_type"] = (
                    type(self.element_type).__name__
                    if field_type == "ScalarField"
                    else type(self.element_type.elem).__name__
                )
                field.create_dataset(
                    name="value", data=self.value, compression="gzip", shuffle=True
                )
        elif file_path.suffix == ".vtu":
            self.to_pyvista().save(file_path)
        else:
            raise NotImplementedError(
                f"The field cannot be saved as a {file_path.suffix} file. "
                "Save it either as an HDF5 (.h5) or a VTK UnstructuredGrid (.vtu) file."
            )

    @abstractmethod
    def plot(self: Self):
        """Plot the field using PyVista."""

    @abstractmethod
    def __str__(self: Self) -> str:
        """Return a concise string representation.

        Returns:
            A human-readable string describing the field.

        """

    @abstractmethod
    def __repr__(self: Self) -> str:
        """Return an unambiguous string representation.

        Returns:
            A string suitable for interactive display.

        """

    def _repr_html_(self: Self):
        """Return a rich HTML representation for Jupyter."""
        field_type = type(self).__name__
        total_cells = 0
        cells_list = (
            self.mesh.cells if isinstance(self.mesh.cells, list) else [self.mesh.cells]
        )
        for cells in cells_list:
            total_cells += cells.connectivity.shape[0]
        max_point = [round(float(self.mesh.points[:, i].max()), 3) for i in range(3)]
        min_point = [round(float(self.mesh.points[:, i].min()), 3) for i in range(3)]
        template = ansh.jinja_env.get_template("field.html.jinja")
        return template.render(
            mesh=self.mesh,
            mesh_max_point=str(max_point),
            mesh_min_point=str(min_point),
            total_cells=total_cells,
            field_type=field_type,
            element_type=type(self.element_type).__name__
            if field_type == "ScalarField"
            else type(self.element_type.elem).__name__,
            field_min=self.value.min(),
            field_max=self.value.max(),
            array_repr=array_repr(self.value),
            uid=str(uuid.uuid4()),
        )

    @abstractmethod
    def __eq__(self: Self, other: object) -> bool:
        """Compare two fields for equality.

        Returns:
            ``True`` if same mesh, element type, and values, otherwise ``False``.

        """

    @abstractmethod
    def __call__(self: Self, point: list | tuple | NDArray) -> NDArray:
        """Evaluate the field at one or more points.

        Args:
            point: A single point (length-3 list/tuple/array) or an array of
                points.

        Returns:
            Interpolated field values at the requested points.

        """

    @classmethod
    def from_file(
        cls,
        file_path: pl.Path | str,
        subregion_tags_key: str | None = None,
        value_key: str | None = None,
    ) -> Self:
        """Load a field from an HDF5 file.

        Args:
            file_path: Path to the file to read.
            subregion_tags_key: Optionally required for ``.vtu`` file.
            value_key: Optionally required for ``.vtu`` file.

        Returns:
            A new field instance of the appropriate type.

        Raises:
            TypeError: If the field type stored in the file does not match the
                calling class.
            NotImplementedError: If the file extension is not supported.

        Notes:
            The ``subregion_tags_key`` is used to point to the cell data, in a ``.vtu``
            file, which stores subregions as integer numbers on the mesh cells.

            The ``value_key`` is used to point to the point data, in a ``.vtu`` file,
            which will be used as the value of the field.

        """
        file_path = file_path if isinstance(file_path, pl.Path) else pl.Path(file_path)
        if file_path.suffix == ".h5":
            mesh = Mesh.from_file(file_path)
            with h5py.File(name=file_path, mode="r") as handle:
                field_type = handle["field"].attrs["type"]
                if cls.__name__ != field_type:
                    raise TypeError(
                        f"Cannot create {cls.__name__} from {field_type} HDF5 file."
                    )
                element_type = eval(f"skfem.{handle['field'].attrs['element_type']}")
                field_value = handle["field/value"][...]
                return cls(mesh, element_type(), field_value)
        elif file_path.suffix == ".vtu":
            mesh = Mesh.from_file(
                file_path=file_path, subregion_tags_key=subregion_tags_key
            )
            pv_grid: pv.UnstructuredGrid = pv.read(file_path)

            if "value" in pv_grid.point_data:
                is_point_data = True
                value_array = pv_grid.point_data["value"]
            elif "value" in pv_grid.cell_data:
                is_point_data = False
                value_array = pv_grid.cell_data["value"]
            elif value_key:
                if value_key in pv_grid.point_data:
                    is_point_data = True
                    value_array = pv_grid.point_data[value_key]
                elif value_key in pv_grid.cell_data:
                    is_point_data = False
                    value_array = pv_grid.cell_data[value_key]
            else:
                raise KeyError("Unable to find key pointing to field data.")

            if value_array.ndim == 2 and cls.__name__ == "ScalarField":
                raise TypeError("Cannot create ScalarField from vector data.")
            elif value_array.ndim == 1 and cls.__name__ == "VectorField":
                raise TypeError("Cannot create VectorField from scalar data.")

            if isinstance(mesh.to_skfem(), skfem.mesh.MeshLine1):
                element = (
                    skfem.ElementLineP1() if is_point_data else skfem.ElementLineP0()
                )
            elif isinstance(mesh.to_skfem(), skfem.mesh.MeshTri1):
                element = (
                    skfem.ElementTriP1() if is_point_data else skfem.ElementTriP0()
                )
            elif isinstance(mesh.to_skfem(), skfem.mesh.MeshTet1):
                element = (
                    skfem.ElementTetP1() if is_point_data else skfem.ElementTetP0()
                )
            elif isinstance(mesh.to_skfem(), skfem.mesh.MeshQuad1):
                element = (
                    skfem.ElementQuad1() if is_point_data else skfem.ElementQuad0()
                )
            elif isinstance(mesh.to_skfem(), skfem.mesh.MeshHex1):
                element = skfem.ElementHex1() if is_point_data else skfem.ElementHex0()
            else:
                raise TypeError(f"Unrecognised mesh type {mesh.to_skfem()}.")

            temp_field = cls(mesh, element, value_array[0])

            compare_points = (
                pv_grid.points if is_point_data else pv_grid.cell_centers().points
            )

            if compare_points.shape == temp_field.dof_locations.shape and np.allclose(
                compare_points, temp_field.dof_locations
            ):
                final_value_array = value_array
            else:
                tree = KDTree(compare_points)
                dist, new_point_indices = tree.query(temp_field.dof_locations)
                if not np.allclose(np.zeros_like(dist), dist):
                    raise RuntimeError("Unable to create new indices dof_locations.")
                final_value_array = value_array[new_point_indices]

            return cls(mesh, element, final_value_array)

        else:
            raise NotImplementedError(
                "The field can be created from "
                "either an HDF5 (.h5) file or a VTK (.vtu). "
                f"Instead a {file_path.suffix} file was given."
            )